"""
Convet log file to csv file
"""
import os
import re
import csv


def save_csv(logs, prefix=""):
    name = prefix + logs['name']
    test_loss = logs['test_loss']
    test_acc = logs['test_acc']
    epochs = logs['epochs']
    train_loss = logs['train_loss']
    train_acc = logs['train_acc']
    learning_rate = logs['lr']
    total = len(epochs)

    print('\t===> write %s.csv' % name)
    with open(name + '.csv', 'w', encoding='utf8', newline='') as f:
        writer = csv.writer(f)

        writer.writerow(
            ['epochs', 'train_loss', 'train_acc', 'lr', 'test_loss', 'test_acc'])

        for i in range(total):
            writer.writerow([epochs[i], train_loss[i], train_acc[i],
                             learning_rate[i], test_loss[i], test_acc[i]])


def deal(log, train_data_size=50000):
    print('\t===> get informations')
    test_loss = []
    test_acc = []
    epochs = []
    train_loss = []
    train_acc = []
    learning_rate = []
    name = None

    regex_train = re.compile('Loss: (.*?) \| Acc: (.*?)% \(\d+/(\d+)\)')
    regex_test = re.compile(
        'Test set: Average loss: (.*), Accuracy: \d+/\d+ \((.*)%\)')
    regex_epoch = re.compile(
        'Epoch: (\d+), lr: (.*?) on model: (.*)')

    for line in log:
        line = line.strip()
        if line.startswith('Loss:'):
            match = regex_train.match(line)
            if match:
                loss = float(match.group(1))
                acc = float(match.group(2))
                num = int(match.group(3))
                if num == train_data_size:
                    train_loss.append(loss)
                    train_acc.append(acc)
        elif line.startswith('Test'):
            match = regex_test.match(line)
            if match:
                loss = float(match.group(1))
                acc = float(match.group(2))
                test_loss.append(loss)
                test_acc.append(acc)
        elif line.startswith('Epoch'):
            match = regex_epoch.match(line)
            if match:
                epoch = int(match.group(1))
                lr = float(match.group(2))
                name = match.group(3)
                epochs.append(epoch)
                learning_rate.append(lr)
    return {'name': name, 'epochs': epochs, 'train_loss': train_loss, 'train_acc': train_acc, 'lr': learning_rate, 'test_loss': test_loss, 'test_acc': test_acc}


def open_log(fname):
    print('===> read log file: ', fname)
    if not os.path.exists(fname):
        print("No such file")
        exit(-1)
    with open(fname, 'r') as f:
        log = f.readlines()
        logs = deal(log)
        save_csv(logs, root)
        print('\t===> finished!')


root = os.path.join(os.getcwd(), 'Experments_log/')
for file in os.listdir(root):
    if file.endswith('.log'):
        open_log(os.path.join(root, file))

# open_log(os.path.join(root, 'VGG_16_full.log'))
# open_log(os.path.join(root, 'VGG_16_full.log'))
# open_log(os.path.join(root, 'VGG_16_tb_0.0.log'))
